import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from collections import Counter
import random
import os
from typing import List, Tuple, Dict

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def main():
    # 读取数据
    data_path = "csv/dlt_lottery_numbers.csv"
    df = pd.read_csv(data_path)

    # 1. 数据预处理
    all_front = []
    for i in range(1, 6):
        all_front.extend(df[f'前区{i}'].astype(int).tolist())

    all_back = []
    for i in range(1, 3):
        all_back.extend(df[f'后区{i}'].astype(int).tolist())

    # 2. 频率统计
    front_counts = Counter(all_front)
    front_freq = pd.DataFrame.from_dict(front_counts, orient='index', columns=['频次']).sort_index()
    front_freq['频率'] = front_freq['频次'] / front_freq['频次'].sum()

    back_counts = Counter(all_back)
    back_freq = pd.DataFrame.from_dict(back_counts, orient='index', columns=['频次']).sort_index()
    back_freq['频率'] = back_freq['频次'] / back_freq['频次'].sum()

    # 3. 可视化分析
    plt.figure(figsize=(18, 15))

    # 前区号码频率分布
    plt.subplot(2, 2, 1)
    sns.barplot(x=front_freq.index, y=front_freq['频次'],
                hue=front_freq.index, palette='viridis', legend=False)
    plt.title('前区号码出现频次 (1-35)')
    plt.xlabel('号码')
    plt.ylabel('出现次数')
    plt.xticks(rotation=45)

    # 前区号码热力图
    plt.subplot(2, 2, 2)
    heatmap_data = np.zeros(35)
    for num, count in front_counts.items():
        heatmap_data[num - 1] = count
    heatmap_data = heatmap_data.reshape(7, 5)  # 7行5列

    xtick_labels = ['1', '6', '11', '16', '21']
    ytick_labels = ['1-5', '6-10', '11-15', '16-20', '21-25', '26-30', '31-35']

    sns.heatmap(heatmap_data, annot=True, fmt='g', cmap='YlOrRd',
                xticklabels=xtick_labels,
                yticklabels=ytick_labels)
    plt.title('前区号码热力图')
    plt.xlabel('号码区间')
    plt.ylabel('号码区间')

    # 后区号码频率分布
    plt.subplot(2, 2, 3)
    sns.barplot(x=back_freq.index, y=back_freq['频次'],
                hue=back_freq.index, palette='coolwarm', legend=False)
    plt.title('后区号码出现频次 (1-12)')
    plt.xlabel('号码')
    plt.ylabel('出现次数')

    # 后区号码热力图
    plt.subplot(2, 2, 4)
    heatmap_data = np.zeros(12)
    for num, count in back_counts.items():
        heatmap_data[num - 1] = count
    heatmap_data = heatmap_data.reshape(3, 4)  # 3行4列

    xtick_labels = ['1', '5', '9', '12']
    ytick_labels = ['1-4', '5-8', '9-12']

    sns.heatmap(heatmap_data, annot=True, fmt='g', cmap='Blues',
                xticklabels=xtick_labels,
                yticklabels=ytick_labels)
    plt.title('后区号码热力图')
    plt.xlabel('号码区间')
    plt.ylabel('号码区间')

    plt.tight_layout()
    plt.savefig('./images/dlt_frequency_analysis.png', dpi=300)
    plt.show()

    # 4. 历史分布规律分析
    print("\n前区号码分析:")
    print(f"最常出现的5个号码: {sorted(front_counts.items(), key=lambda x: x[1], reverse=True)[:5]}")
    print(f"最少出现的5个号码: {sorted(front_counts.items(), key=lambda x: x[1])[:5]}")

    print("\n后区号码分析:")
    print(f"最常出现的3个号码: {sorted(back_counts.items(), key=lambda x: x[1], reverse=True)[:3]}")
    print(f"最少出现的3个号码: {sorted(back_counts.items(), key=lambda x: x[1])[:3]}")

    # 计算号码区间分布
    def calculate_interval_distribution(numbers: List[int], intervals: List[Tuple[int, int]]) -> Dict[str, int]:
        """计算号码在指定区间内的分布"""
        dist = {}
        for num in numbers:
            for start, end in intervals:
                if start <= num <= end:
                    interval_name = f"{start}-{end}"
                    dist[interval_name] = dist.get(interval_name, 0) + 1
                    break
        return dist

    # 前区区间划分
    front_intervals = [(1, 5), (6, 10), (11, 15), (16, 20), (21, 25), (26, 30), (31, 35)]
    front_interval_dist = calculate_interval_distribution(all_front, front_intervals)

    # 后区区间划分
    back_intervals = [(1, 4), (5, 8), (9, 12)]
    back_interval_dist = calculate_interval_distribution(all_back, back_intervals)

    print("\n前区号码区间分布:")
    for interval, count in sorted(front_interval_dist.items()):
        print(f"{interval}: {count}次 ({count / len(all_front) * 100:.1f}%)")

    print("\n后区号码区间分布:")
    for interval, count in sorted(back_interval_dist.items()):
        print(f"{interval}: {count}次 ({count / len(all_back) * 100:.1f}%)")

    # 5. 号码推荐策略
    def generate_recommendation(strategy: str = 'balanced') -> Tuple[List[int], List[int]]:
        """
        生成大乐透号码推荐
        策略选项:
        - 'hot': 只选择高频号码
        - 'cold': 只选择低频号码
        - 'balanced': 平衡选择高低频号码
        - 'random': 完全随机选择
        """
        # 前区选择 (1-35选5)
        front_nums = list(range(1, 36))

        # 后区选择 (1-12选2)
        back_nums = list(range(1, 13))

        if strategy == 'hot':
            # 选择前区最高频的5个号码
            front_selected = sorted(front_counts.items(), key=lambda x: x[1], reverse=True)[:5]
            front_selected = [num for num, _ in front_selected]

            # 选择后区最高频的2个号码
            back_selected = sorted(back_counts.items(), key=lambda x: x[1], reverse=True)[:2]
            back_selected = [num for num, _ in back_selected]

        elif strategy == 'cold':
            # 选择前区最低频的5个号码
            front_selected = sorted(front_counts.items(), key=lambda x: x[1])[:5]
            front_selected = [num for num, _ in front_selected]

            # 选择后区最低频的2个号码
            back_selected = sorted(back_counts.items(), key=lambda x: x[1])[:2]
            back_selected = [num for num, _ in back_selected]

        elif strategy == 'balanced':
            # 平衡策略：混合高低频号码
            # 前区：3个高频 + 2个低频
            hot_front = sorted(front_counts.items(), key=lambda x: x[1], reverse=True)[:10]  # 取前10高频
            cold_front = sorted(front_counts.items(), key=lambda x: x[1])[:10]  # 取前10低频

            # 随机选择3个高频和2个低频
            front_selected = [num for num, _ in random.sample(hot_front, 3)] + [num for num, _ in
                                                                                random.sample(cold_front, 2)]

            # 后区：1个高频 + 1个低频
            hot_back = sorted(back_counts.items(), key=lambda x: x[1], reverse=True)[:5]  # 取前5高频
            cold_back = sorted(back_counts.items(), key=lambda x: x[1])[:5]  # 取前5低频

            back_selected = [num for num, _ in random.sample(hot_back, 1)] + [num for num, _ in random.sample(cold_back, 1)]

        else:  # random
            # 完全随机选择
            front_selected = random.sample(front_nums, 5)
            back_selected = random.sample(back_nums, 2)

        # 排序
        front_selected.sort()
        back_selected.sort()

        return front_selected, back_selected

    # 6. 生成推荐号码
    print("\n基于历史数据分析的号码推荐:")
    strategies = ['hot', 'cold', 'balanced', 'random']
    recommendations = []

    for strategy in strategies:
        front, back = generate_recommendation(strategy)
        recommendations.append({
            '策略': strategy,
            '前区': front,
            '后区': back,
            '组合': f"{' '.join(map(str, front))} + {' '.join(map(str, back))}"
        })
        print(f"{strategy.upper()}策略: {' '.join(map(str, front))} + {' '.join(map(str, back))}")

    # 保存推荐结果
    recommend_df = pd.DataFrame(recommendations)
    print("\n推荐号码汇总:")
    print(recommend_df)

    # 保存到文件
    os.makedirs('../lottery/recommendations', exist_ok=True)
    recommend_df.to_csv('../lottery/recommendations/dlt_recommendations.csv', index=False)
    print("\n推荐结果已保存到: ../lottery/recommendations/dlt_recommendations.csv")

if __name__ == "__main__":
    main()